#include <cstdio>
#include <unistd.h>
#include "Master.h"
#include "Block.h"
#include "BlockStorage.h"
#include "Signals.h"
#include "mpi.h"
#include "Slave.h"


class BlockTest : public Block {
    int data;
public:
    BlockTest(int id) {
        this -> id = id;
        data = id;
        fprintf(stderr,"Creating block nr %d\n", id);
    }
    void dispose() {
        fprintf(stderr,"Deleting block nr %d\n", id);
    }
    
protected:
    void* getRawDataBuffer() { return &data; }
    int getRawDataSize() { return 1; }
    MPI_Datatype getDataType() { return MPI_INT; }
};

class BlockStorageTest : public BlockStorage {
protected:
    Block* createBlock(int block_id) {
        return new BlockTest(block_id);
    }
};

class BlockStorageTest2 : public BlockStorageTest {
public:
    BlockStorageTest2(int k) {
        for(int i = k; i < k + 5; ++i) {
            addToStorage(new BlockTest(i));
        }
    }
};

class MasterTest : public Master {

public:
    MasterTest() {
        blocks_count = 10;
        iterations_count = 2;
        for(int i = 0; i < 5; ++i) {
            markBlockAsDone(i, 1);
        }
        for(int i = 5; i < 10; ++i) {
            markBlockAsDone(i, 2);
        }
    }
};

class SlaveTest : public Slave {
    int id;
public:
    SlaveTest(int id) {
        storage = new BlockStorageTest2((id - 1) * 5);
        this -> id = id;
    }
    ~SlaveTest() {
        delete storage;
    }
protected:
    void updateBlock() {
        Block* block = storage -> getBlock(working_block_id);
        sleep(id);
        printf("%d: Done working on %d block...\n", id, block -> getId());
    }
};

int main(int argc, char* argv[]) {
    MPI_Init(&argc, &argv);
    int id, nothing;
    MPI_Comm_rank(MPI_COMM_WORLD, &id);
    
    if(id == 0) {
        MasterTest t;
        #pragma omp parallel sections
        {
            #pragma omp section
            {
                t.serveBlockInfo();
                printf("Serving block info done\n");
            }
            #pragma omp section
            {
                t.run();
                printf("Master ends work\n");
            }
        }
    }
    else{
        SlaveTest t(id);
        #pragma omp parallel sections
        {
            #pragma omp section
            {
                t.serveBlocks();
                printf("%d: Serving blocks done\n", id);
            }
            #pragma omp section
            {
                t.run();
                printf("%d: Slave work done\n", id);
            }
        }
    }
    
    MPI_Finalize();
    return 0;
}
